#include <iostream>
#include <cuda_runtime.h>
#include "cublas_v2.h"
#include "../common/common.h"
#define DEVICENUM 0
#define M 5 // 向量的长度

int main(int argc, char **argv)
{
    cudaSetDevice(DEVICENUM);

    // 声明主机端矩阵A
    float *A = (float *)malloc(sizeof(float) * M * M);
    float *vector_dtoh_res = (float *)malloc(sizeof(float) * M);
    initialDataIncreasing(A, M * M);
    memset(vector_dtoh_res, 0, sizeof(float) * M);
    printMatrix(A, M, M, "A");

    // 声明设备端向量dX，设备端向量dY,并出申请内存
    float *dX, *dY;
    CHECK(cudaMalloc((void **)&dX, sizeof(float) * M));
    CHECK(cudaMalloc((void **)&dY, sizeof(float) * M));

    cublasHandle_t handle = 0; // 创建句柄，这里要进行向量计算了，所以要用到
    CHECK_CUBLAS(cublasCreate(&handle));

    // 1. 矩阵A中的一行到向量dY的复制
    CHECK_CUBLAS(cublasSetVector(M, sizeof(float), A, 1, dY, 1));
    CHECK_CUBLAS(cublasGetVector(M, sizeof(float), dY, 1, vector_dtoh_res, 1));
    printVector(vector_dtoh_res, M, "dY");

    // 2. 矩阵A中的一列到向量dX的复制
    CHECK_CUBLAS(cublasSetVector(M, sizeof(float), A, M, dX, 1));
    CHECK_CUBLAS(cublasGetVector(M, sizeof(float), dX, 1, vector_dtoh_res, 1));
    printVector(vector_dtoh_res, M, "dX");

    // 3. 使用cublasSdot完成点积计算
    float *dot_res = (float *)malloc(sizeof(float));
    memset(dot_res, 0, sizeof(float)); // 将dot_res置为0
    CHECK_CUBLAS(cublasSdot(handle, M, dX, 1, dY, 1, dot_res));
    std::cout << "dX * dY = " << *dot_res << std::endl;

    // 4. 也可使用较高级的用法，直接从矩阵中取行或列进行操作
    float *dA; // 创建设备端矩阵
    CHECK(cudaMalloc((void **)&dA, sizeof(float) * M * M));
    CHECK_CUBLAS(cublasSetMatrix(M, M, sizeof(float), A, M, dA, M));

    memset(dot_res, 0, sizeof(float)); // 将dot_res置为0
    CHECK_CUBLAS(cublasSdot(handle, M, dA, M, dA, 1, dot_res));
    std::cout << "dA_col * dA_row = " << *dot_res << std::endl;
    cudaFree(dA);

    // 释放内存
    free(A);
    free(dot_res);
    free(vector_dtoh_res);
    cudaFree(dX);
    cudaFree(dY);

    CHECK_CUBLAS(cublasDestroy(handle));

    return 0;
}